import numpy as np

class NaiveBayes:
    """朴素贝叶斯分类器实现"""
    
    def __init__(self):
        # 存储每个类别的先验概率
        self.prior_prob = {}
        # 存储每个特征在每个类别下的条件概率
        self.cond_prob = {}
        # 存储所有类别
        self.classes = None
        # 存储每个特征的可能取值
        self.features = None
        # 拉普拉斯平滑参数
        self.laplace = 1
    
    def fit(self, X, y):
        """
        训练朴素贝叶斯分类器
        X: 特征数据，二维列表或数组
        y: 标签数据，一维列表或数组
        """
        # 转换为numpy数组便于操作
        X = np.array(X)
        y = np.array(y)
        
        # 获取样本数量和特征数量
        n_samples, n_features = X.shape
        
        # 获取所有类别
        self.classes = np.unique(y)
        n_classes = len(self.classes)
        
        # 存储每个特征的可能取值
        self.features = [np.unique(X[:, i]) for i in range(n_features)]
        
        # 计算先验概率 P(class)
        for c in self.classes:
            self.prior_prob[c] = np.sum(y == c) / n_samples
        
        # 计算条件概率 P(feature|class)
        for c in self.classes:
            # 获取属于类别c的所有样本
            X_c = X[y == c]
            n_samples_c = X_c.shape[0]
            
            self.cond_prob[c] = []
            for i in range(n_features):
                # 计算特征i在类别c下的条件概率
                feature_probs = {}
                # 获取特征i的所有可能取值
                feature_values = self.features[i]
                
                for value in feature_values:
                    # 计算该取值出现的次数
                    count = np.sum(X_c[:, i] == value)
                    # 应用拉普拉斯平滑
                    prob = (count + self.laplace) / (n_samples_c + self.laplace * len(feature_values))
                    feature_probs[value] = prob
                
                self.cond_prob[c].append(feature_probs)
    
    def predict(self, X):
        """
        预测样本类别
        X: 特征数据，二维列表或数组
        返回预测结果列表
        """
        X = np.array(X)
        predictions = []
        
        for sample in X:
            # 计算每个类别的后验概率（取对数避免下溢）
            post_probs = {}
            for c in self.classes:
                # 初始化为先验概率的对数
                log_prob = np.log(self.prior_prob[c])
                
                # 累加每个特征的条件概率的对数
                for i, value in enumerate(sample):
                    # 如果特征值从未见过，赋予一个很小的概率
                    if value not in self.cond_prob[c][i]:
                        log_prob += np.log(self.laplace / (len(self.features[i]) * self.laplace))
                    else:
                        log_prob += np.log(self.cond_prob[c][i][value])
                
                post_probs[c] = log_prob
            
            # 选择后验概率最大的类别作为预测结果
            prediction = max(post_probs, key=post_probs.get)
            predictions.append(prediction)
        
        return predictions

# 示例用法
if __name__ == "__main__":
    # 天气数据集：[outlook, temperature, humidity, wind]
    X = [
        ['sunny', 'hot', 'high', 'weak'],
        ['sunny', 'hot', 'high', 'strong'],
        ['overcast', 'hot', 'high', 'weak'],
        ['rain', 'mild', 'high', 'weak'],
        ['rain', 'cool', 'normal', 'weak'],
        ['rain', 'cool', 'normal', 'strong'],
        ['overcast', 'cool', 'normal', 'strong'],
        ['sunny', 'mild', 'high', 'weak'],
        ['sunny', 'cool', 'normal', 'weak'],
        ['rain', 'mild', 'normal', 'weak'],
        ['sunny', 'mild', 'normal', 'strong'],
        ['overcast', 'mild', 'high', 'strong'],
        ['overcast', 'hot', 'normal', 'weak'],
        ['rain', 'mild', 'high', 'strong']
    ]
    
    # 是否适合打网球
    y = [
        'no', 'no', 'yes', 'yes', 'yes', 'no', 'yes', 'no',
        'yes', 'yes', 'yes', 'yes', 'yes', 'no'
    ]
    
    # 创建并训练模型
    nb = NaiveBayes()
    nb.fit(X, y)
    
    # 测试预测
    test_samples = [
        ['sunny', 'hot', 'high', 'weak'],
        ['overcast', 'mild', 'normal', 'strong'],
        ['rain', 'cool', 'high', 'weak']
    ]
    
    predictions = nb.predict(test_samples)
    print("测试样本预测结果：", predictions)
